package com.nit.netty.manager;

import com.nit.netty.bean.invoker.NettyInterfaceInvoker;
import com.nit.netty.sender.RequestSender;
import com.nit.netty.zookeeper.client.ZookeeperClient;
import io.netty.channel.Channel;
import lombok.extern.slf4j.Slf4j;
import net.jodah.expiringmap.ExpirationPolicy;
import net.jodah.expiringmap.ExpiringMap;
import org.apache.zookeeper.KeeperException;

import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Author: whb
 * @Description: 一个channel的管理者, 获取channel和新建channel
 * @CreateTime: 2020/8/21 11:27
 */
@Slf4j
public class NettyChannelManager {
    private final Map<String, List<SocketAddress>> serviceAddressMap;
    private final ExpiringMap<SocketAddress, Channel> channels;
    private final Map<String, AtomicInteger> roundRobinMap;
    private final RequestSender sender;
    private final ZookeeperClient zkClient;

    public NettyChannelManager(RequestSender sender, ZookeeperClient zkClient) {
        this.serviceAddressMap = new ConcurrentHashMap<>();
        this.roundRobinMap = new ConcurrentHashMap<>();
        this.sender = sender;
        this.zkClient = zkClient;
        // 给zk设置manager, 事件监听调用refresh
        zkClient.setManager(this);
        channels = ExpiringMap.builder()
                .maxSize(100)
                .expiration(10, TimeUnit.MINUTES) // 十分钟没有通信就会关闭channel
                .expirationPolicy(ExpirationPolicy.ACCESSED)
                .variableExpiration()
                .expirationListener(this::removeExpiredChannel)
                .build();
        NettyInterfaceInvoker.setSender(this.sender);
    }
    // 这里只关闭channel的连接, 不会去清理key
    private void removeExpiredChannel(SocketAddress key, Channel channel) {
        if (channel != null) {
            channel.close();
        }
        channels.put(key, channel);
    }

    /**
     * @Description: 轮询获取一个channel
     */
    public synchronized Channel take(String serviceName) {
        List<SocketAddress> socketAddressList = serviceAddressMap.get(serviceName);
        // 服务名称下不存在服务端连接则返回null
        if (socketAddressList != null && socketAddressList.size() > 0) {
            int size = channels.size();
            AtomicInteger roundRobin = roundRobinMap.get(serviceName);
            int index = (roundRobin.getAndAdd(1) + size) % size;
            SocketAddress socketAddress = socketAddressList.get(index);
            Channel channel = channels.get(socketAddress);
            if (channel != null && channel.isOpen()) {
                // 重置超时时间 (这里可能用put会更好, 极端情况下key失效了就无法reset了
                channels.resetExpiration(socketAddress);
                return channel;
            } else if(channel != null && !channel.isRegistered()) {
                // channel超时后, 关闭连接, 这里需要重启连接
                try {
                    // 重新获取与服务端的连接
                    channel = connectServerNode(socketAddress);
                } catch(IllegalArgumentException e) {
                    // 如果获取失败就说明是服务端主动关闭的, 或者服务端网络异常
                    log.error("客户端获取channel失败, 移除失效channel后重试");
                    // 这里会重新递归, 移除了失效channel之后的channels会返回一个有效的channel
                    return take(serviceName);
                }
                return channel;
            }
        }
        return null;
    }

    public synchronized void refresh() {
        try {
            // refresh目的是重新从zookeeper上拉取最新服务, 所以原有服务map清空
            serviceAddressMap.clear();
            List<String> serviceList = zkClient.getAllService();
            // 若zookeeper上无任何服务, 则直接将channel清空
            if (serviceList == null || serviceList.size() == 0) {
                for (final Channel channel : channels.values()) {
                    channel.close();
                }
                channels.clear();
                log.info("zookeeper上无任何服务");
                return;
            }
            // 获取最新服务
            refreshService(serviceList);

            // 检查并且移除已经无用或者失效的channel
            checkAndRemoveChannel();

            log.info("服务refresh成功");
        } catch (KeeperException | InterruptedException e) {
            e.printStackTrace();
        }
    }

    /**
     * @Description: 获取zookeeper上的最新服务
     */
    private void refreshService(List<String> serviceList) throws KeeperException, InterruptedException {
        // 遍历服务列表
        for (String serviceName : serviceList) {
            // 获取服务下的所有地址节点
            List<String> addressList = zkClient.getServiceByName(serviceName);
            // 遍历地址节点, 获得真正的地址
            List<SocketAddress> socketAddressList = new ArrayList<>();
            serviceAddressMap.put(serviceName, socketAddressList);
            roundRobinMap.put(serviceName, new AtomicInteger(0));
            for (String address : addressList) {
                String path = serviceName + "/" + address;
                String[] hostAndPort = zkClient.getData(path).split(":");
                if (hostAndPort.length == 2) {
                    String host = hostAndPort[0];
                    int port = Integer.parseInt(hostAndPort[1]);
                    final SocketAddress remotePeer = new InetSocketAddress(host, port);
                    socketAddressList.add(remotePeer);
                    checkAndBuildChannel(remotePeer);
                }
            }
        }
    }

    /**
     * @Description: 检查并移除channel
     */
    private void checkAndRemoveChannel() {
        // 创建一个需要移除的List
        List<SocketAddress> removeList = new ArrayList<>();
        // 创建一个连接地址的最全集
        Set<SocketAddress> allAddress = new HashSet<>();
        // 获取一个服务的所有连接, 此时已经是最新状态
        for (List<SocketAddress> socketAddressSet : serviceAddressMap.values()) {
            allAddress.addAll(socketAddressSet);
        }
        // 获取获取所有channel对应的地址
        // expiringMap的keySet有bug 会无限循环, 所以这里只能曲线救国
        // 已经提了issue, 就看作者维不维护了
        int count = 0;
        for (SocketAddress socketAddress : channels.keySet()) {
            count ++;
            Channel channel = channels.get(socketAddress);
            // 地址不为空并且在所有服务中都无此地址
            if (socketAddress != null && !allAddress.contains(socketAddress)) {
                // 做关闭channel操作
                if (channel != null) {
                    channel.close();
                }
                channel = null; // help GC
                // foreach的时候迭代删除会出错
                // 可以使用iterator解决
                removeList.add(socketAddress);
            }
            if(count == channels.size()) {
                break;
            }
        }
        for (SocketAddress socketAddress : removeList) {
            channels.remove(socketAddress);
        }
    }

    /**
     * @Description: 检查并构建channel
     */
    private void checkAndBuildChannel(SocketAddress serverNodeAddress) {
        // 一个地址对应一个channel, 如果channel存在就不会重新连接
        Channel channel = channels.get(serverNodeAddress);
        if (channel != null && channel.isOpen()) {
            log.info("当前服务节点已存在,无需重新连接.{}", serverNodeAddress);
        } else {
            // 此时channel已经关闭 直接移除新建
            if (channel != null) {
                channel = null; // help GC
                channels.remove(serverNodeAddress);
            }
            try {
                connectServerNode(serverNodeAddress);
            } catch(IllegalArgumentException e) {
                log.error("refresh的时候channel已被关闭",e);
            }
        }
    }

    /**
     * @Description: refresh的时候重新构建channel放入map
     */
    private Channel connectServerNode(SocketAddress address) throws IllegalArgumentException {
        try {
            Channel channel = sender.connect(address);
            // 这里说明是服务器主动的断开连接, 我们需要去除这个端点
            if(null == channel) {
                remove(address);
                // 移除后抛出异常
                throw new IllegalArgumentException();
            }
            channels.put(address, channel);
            return channel;
        } catch (InterruptedException e) {
            e.printStackTrace();
            log.info("未能连接到服务器.{}", address);
        }
        return null;
    }

    public synchronized void remove(SocketAddress address) {
        channels.remove(address);
        for (List<SocketAddress> list : serviceAddressMap.values()) {
            list.remove(address);
        }
    }
}

